import nvtx

class NVTXModuleDecorator():
    '''
    NVTX wrapper for forward method, with arguments
    '''
    def __init__(self, color='green', desc=None):
        self.desc = desc
        self.color = color
    
    def __call__(self, cls):
        # 接收被装饰的类
        class WrappedClass(cls):
            desc = self.desc
            color = self.color
            
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.desc = self.desc or cls.__name__

            def forward(self, *args, **kwargs):
                with nvtx.annotate(message=self.desc, color=self.color):
                    return super().forward(*args, **kwargs)
        return WrappedClass

def nvtx_forward_decorator(cls):
    '''
    nvtx wrapper for forward method, without arguments
    '''
    # record original func
    origin_forward = getattr(cls, 'forward', None)

    if origin_forward and callable(origin_forward):
        def nvtx_forward_wrapper(*args, **kwargs):
            with torch.annotate(message=cls.__name__ + '.forward', color='green'):
                return origin_forward(*args, **kwargs)
        setattr(cls, 'forward', nvtx_forward_wrapper)
    return cls
